#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
主窗口实现
"""

import os
import sys
import json
import cv2
from pathlib import Path
from PyQt5.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
    QAction, QMenu, QToolBar, QStatusBar, QFileDialog, 
    QMessageBox, QDockWidget, QLabel, QInputDialog, QProgressDialog
)
from PyQt5.QtCore import Qt, QSize, QRectF
from PyQt5.QtGui import QIcon, QColor
from PyQt5.QtWidgets import QApplication
import numpy as np

from core.dataset import DatasetManager
from core.training import ModelTrainer
from gui.widgets import ImageCanvas, LabelManager, AnnotationToolbar, DatasetBrowser, TrainingMonitor
from gui.windows.dialogs import NewDatasetDialog, AddLabelDialog, DatasetSelectDialog, ROIDialog

class MainWindow(QMainWindow):
    """主窗口类"""
    
    def __init__(self, config, logger, parent=None):
        """
        初始化主窗口
        
        Args:
            config (dict): 应用配置
            logger (logging.Logger): 日志器
            parent (QWidget, optional): 父窗口
        """
        super().__init__(parent)
        self.config = config
        self.logger = logger
        
        # 存储当前数据集信息
        self.current_dataset = None
        
        # 创建数据集管理器
        self.dataset_manager = DatasetManager(config['paths']['datasets'], logger)
        
        # 设置窗口属性
        self.setWindowTitle(config['app']['name'])
        self.resize(1280, 800)
        
        # 初始化UI
        self._setup_ui()
        self._create_actions()
        self._create_menus()
        self._create_toolbars()
        self._create_statusbar()
        self._create_connections()
        
        self.logger.info("主窗口初始化完成")
    
    def _setup_ui(self):
        """设置界面布局"""
        # 主部件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        
        # 主布局
        self.main_layout = QVBoxLayout(self.central_widget)
        
        # 提示文本
        self.welcome_label = QLabel("欢迎使用视觉模型训练软件\n请创建或打开一个数据集开始工作")
        self.welcome_label.setAlignment(Qt.AlignCenter)
        self.welcome_label.setStyleSheet("font-size: 18px; color: #808080;")
        
        self.main_layout.addWidget(self.welcome_label)
        
        # 创建画布（初始隐藏）
        self.canvas = ImageCanvas()
        self.canvas.setVisible(False)
        self.main_layout.addWidget(self.canvas)
        
        # 设置停靠窗口
        self._setup_dock_widgets()
    
    def _setup_dock_widgets(self):
        """设置停靠窗口"""
        # 数据集浏览器（左侧）
        self.dataset_dock = QDockWidget("数据集浏览器", self)
        self.dataset_dock.setAllowedAreas(Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea)
        
        # 创建数据集浏览器组件
        self.dataset_browser = DatasetBrowser(self)
        self.dataset_browser.image_selected.connect(self.on_image_selected)
        self.dataset_browser.image_double_clicked.connect(self.on_image_double_clicked)
        self.dataset_browser.image_deleted.connect(self.on_image_deleted)
        self.dataset_dock.setWidget(self.dataset_browser)
        
        self.addDockWidget(Qt.LeftDockWidgetArea, self.dataset_dock)
        
        # 标签管理器（右侧）
        self.label_dock = QDockWidget("标签管理", self)
        self.label_dock.setAllowedAreas(Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea)
        
        # 创建标签管理器组件
        self.label_manager = LabelManager(self)
        self.label_dock.setWidget(self.label_manager)
        
        self.addDockWidget(Qt.RightDockWidgetArea, self.label_dock)
        
        # 训练监控器（底部）
        self.training_dock = QDockWidget("模型训练", self)
        self.training_dock.setObjectName("training_dock")
        self.training_dock.setAllowedAreas(Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea)
        
        self.training_monitor = TrainingMonitor()
        self.training_monitor.start_training.connect(self.on_start_training)
        self.training_monitor.stop_training.connect(self.on_stop_training)
        self.training_monitor.convert_to_annotation.connect(self.on_convert_to_annotation)
        
        self.training_dock.setWidget(self.training_monitor)
        self.addDockWidget(Qt.RightDockWidgetArea, self.training_dock)
        self.training_dock.setVisible(False)  # 初始隐藏
        
        # 初始隐藏这些面板，等打开数据集后显示
        self.dataset_dock.setVisible(False)
        self.label_dock.setVisible(False)
        self.training_dock.setVisible(False)
    
    def _create_actions(self):
        """创建动作"""
        # 文件菜单动作
        self.new_dataset_action = QAction("新建数据集", self)
        self.new_dataset_action.setShortcut("Ctrl+N")
        
        self.open_dataset_action = QAction("打开数据集", self)
        self.open_dataset_action.setShortcut("Ctrl+O")
        
        self.import_images_action = QAction("导入图像...", self)
        self.import_images_action.setShortcut("Ctrl+I")
        self.import_images_action.setEnabled(False)  # 需要先打开数据集
        
        self.import_folder_action = QAction("导入文件夹...", self)
        self.import_folder_action.setShortcut("Ctrl+Shift+I")
        self.import_folder_action.setEnabled(False)  # 需要先打开数据集
        
        self.exit_action = QAction("退出", self)
        self.exit_action.setShortcut("Ctrl+Q")
        
        # 编辑菜单动作
        self.roi_action = QAction("设置ROI区域", self)
        self.roi_action.setEnabled(False)
        
        # 标注菜单动作
        self.create_label_action = QAction("创建标签", self)
        self.create_label_action.setEnabled(False)
        
        self.manage_labels_action = QAction("管理标签", self)
        self.manage_labels_action.setEnabled(False)
        
        # 训练菜单动作
        self.start_training_action = QAction("开始训练", self)
        self.start_training_action.setEnabled(False)
        
        self.export_model_action = QAction("导出模型", self)
        self.export_model_action.setEnabled(False)
    
    def _create_menus(self):
        """创建菜单"""
        # 主菜单
        self.menu_bar = self.menuBar()
        
        # 文件菜单
        self.file_menu = self.menu_bar.addMenu("文件")
        self.file_menu.addAction(self.new_dataset_action)
        self.file_menu.addAction(self.open_dataset_action)
        self.file_menu.addSeparator()
        self.file_menu.addAction(self.import_images_action)
        self.file_menu.addAction(self.import_folder_action)
        self.file_menu.addSeparator()
        self.file_menu.addAction(self.exit_action)
        
        # 编辑菜单
        self.edit_menu = self.menu_bar.addMenu("编辑")
        self.edit_menu.addAction(self.roi_action)
        
        # 标注菜单
        self.annotation_menu = self.menu_bar.addMenu("标注")
        self.annotation_menu.addAction(self.create_label_action)
        self.annotation_menu.addAction(self.manage_labels_action)
        
        # 训练菜单
        self.training_menu = self.menu_bar.addMenu("训练")
        self.training_menu.addAction(self.start_training_action)
        self.training_menu.addAction(self.export_model_action)
        
        # 视图菜单
        self.view_menu = self.menu_bar.addMenu("视图")
        self.view_menu.addAction(self.dataset_dock.toggleViewAction())
        self.view_menu.addAction(self.label_dock.toggleViewAction())
        self.view_menu.addAction(self.training_dock.toggleViewAction())
        
        # 帮助菜单
        self.help_menu = self.menu_bar.addMenu("帮助")
        self.help_menu.addAction("关于")
    
    def _create_toolbars(self):
        """创建工具栏"""
        # 主工具栏
        self.main_toolbar = QToolBar("主工具栏")
        self.main_toolbar.setIconSize(QSize(32, 32))
        self.addToolBar(self.main_toolbar)
        
        # 添加动作
        self.main_toolbar.addAction(self.new_dataset_action)
        self.main_toolbar.addAction(self.open_dataset_action)
        self.main_toolbar.addSeparator()
        self.main_toolbar.addAction(self.import_images_action)
        
        # 标注工具栏
        self.annotation_toolbar_widget = AnnotationToolbar()
        self.annotation_toolbar_widget.tool_changed.connect(self.on_annotation_tool_changed)
        self.annotation_toolbar_widget.color_changed.connect(self.on_annotation_color_changed)
        self.annotation_toolbar_widget.class_changed.connect(self.on_annotation_class_changed)
        
        self.annotation_toolbar = QToolBar("标注工具")
        self.annotation_toolbar.setIconSize(QSize(32, 32))
        self.annotation_toolbar.addWidget(self.annotation_toolbar_widget)
        self.addToolBar(Qt.LeftToolBarArea, self.annotation_toolbar)
        self.annotation_toolbar.setVisible(False)
    
    def _create_statusbar(self):
        """创建状态栏"""
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        
        # 添加永久消息
        self.status_label = QLabel("就绪")
        self.status_bar.addPermanentWidget(self.status_label)
    
    def _create_connections(self):
        """创建信号连接"""
        # 文件操作
        self.new_dataset_action.triggered.connect(self.on_new_dataset)
        self.open_dataset_action.triggered.connect(self.on_open_dataset)
        self.import_images_action.triggered.connect(self.on_import_images)
        self.import_folder_action.triggered.connect(self.on_import_folder)
        self.exit_action.triggered.connect(self.close)
        
        # 标签管理器连接
        self.create_label_action.triggered.connect(self.on_create_label)
        self.manage_labels_action.triggered.connect(self.on_manage_labels)
        self.label_manager.add_button.clicked.connect(self.on_create_label)
        self.label_manager.label_selected.connect(self.on_label_selected)
        self.label_manager.label_color_changed.connect(self.on_label_color_changed)
        self.label_manager.label_removed.connect(self.on_label_removed)
        self.label_manager.label_renamed.connect(self.on_label_renamed)
        
        # 编辑相关
        self.roi_action.triggered.connect(self.on_set_roi)
        
        # 标注相关
        self.canvas.annotation_added.connect(self.on_annotation_added)
        self.canvas.annotation_deleted.connect(self.on_annotation_deleted)
        
        # 训练相关
        self.start_training_action.triggered.connect(self.on_start_training_action)
        self.export_model_action.triggered.connect(self.on_export_model)
    
    def on_new_dataset(self):
        """新建数据集响应函数"""
        self.logger.info("用户请求新建数据集")
        
        # 创建并显示新建数据集对话框
        dialog = NewDatasetDialog(self.config, self)
        result = dialog.exec_()
        
        if result == NewDatasetDialog.Accepted:
            # 获取对话框结果
            dataset_info = dialog.get_result()
            if dataset_info:
                # 创建数据集
                dataset = self.dataset_manager.create_dataset(
                    dataset_info["name"],
                    dataset_info["algorithm_type"]
                )
                
                if dataset:
                    self.current_dataset = dataset
                    self._update_ui_for_dataset()
                    self.status_label.setText(f"当前数据集: {dataset['name']}")
                    self.logger.info(f"已创建数据集: {dataset['name']}")
                else:
                    QMessageBox.warning(self, "创建失败", "无法创建数据集，请查看日志获取详细信息")
    
    def on_open_dataset(self):
        """打开数据集响应函数"""
        self.logger.info("用户请求打开数据集")
        
        # 获取数据集列表
        datasets = self.dataset_manager.get_datasets_list()
        
        if not datasets:
            QMessageBox.information(self, "打开数据集", "没有找到可用的数据集，请先创建一个新的数据集。")
            return
        
        # 显示数据集选择对话框
        dialog = DatasetSelectDialog(datasets, self)
        result = dialog.exec_()
        
        if result == DatasetSelectDialog.Accepted:
            # 获取选中的数据集
            selected_dataset = dialog.get_selected_dataset()
            if selected_dataset:
                # 加载数据集
                dataset = self.dataset_manager.load_dataset(selected_dataset["path"])
                
                if dataset:
                    self.current_dataset = dataset
                    self._update_ui_for_dataset()
                    self.status_label.setText(f"当前数据集: {dataset['name']}")
                    self.logger.info(f"已打开数据集: {dataset['name']}")
                else:
                    QMessageBox.warning(self, "打开失败", "无法打开数据集，请查看日志获取详细信息")
    
    def on_import_images(self):
        """导入图像响应函数"""
        self.logger.info("用户请求导入图像")
        
        if not self.current_dataset:
            QMessageBox.warning(self, "导入错误", "请先打开或创建一个数据集")
            return
        
        # 显示文件选择对话框
        file_dialog = QFileDialog(self)
        file_dialog.setFileMode(QFileDialog.ExistingFiles)
        file_dialog.setNameFilter("图像文件 (*.jpg *.jpeg *.png *.bmp *.tif *.tiff)")
        
        if file_dialog.exec_():
            file_paths = file_dialog.selectedFiles()
            if file_paths:
                # 导入图像
                count = self.dataset_manager.add_images(file_paths)
                if count > 0:
                    QMessageBox.information(
                        self, 
                        "导入成功", 
                        f"已成功导入{count}张图像到数据集"
                    )
                    # 更新界面
                    self._update_image_list()
                else:
                    QMessageBox.warning(self, "导入失败", "没有导入任何图像，请查看日志获取详细信息")
    
    def on_import_folder(self):
        """导入文件夹响应函数"""
        self.logger.info("用户请求导入文件夹")
        
        if not self.current_dataset:
            QMessageBox.warning(self, "导入错误", "请先打开或创建一个数据集")
            return
        
        # 显示文件夹选择对话框
        folder_dialog = QFileDialog(self)
        folder_dialog.setFileMode(QFileDialog.Directory)
        folder_dialog.setOption(QFileDialog.ShowDirsOnly, True)
        
        if folder_dialog.exec_():
            selected_dirs = folder_dialog.selectedFiles()
            if selected_dirs:
                folder_path = selected_dirs[0]
                
                # 获取文件夹中的所有图像
                image_files = []
                
                # 支持的图像格式
                valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']
                
                # 遍历文件夹
                for root, dirs, files in os.walk(folder_path):
                    for file in files:
                        ext = os.path.splitext(file.lower())[1]
                        if ext in valid_extensions:
                            image_files.append(os.path.join(root, file))
                
                if not image_files:
                    QMessageBox.warning(self, "导入失败", "所选文件夹中没有找到支持的图像文件")
                    return
                
                # 询问是否自动标注
                auto_label = False
                if self.current_dataset["algorithm_type"] == "classification":
                    reply = QMessageBox.question(
                        self, 
                        "自动标注",
                        f"检测到{len(image_files)}个图像文件。是否要根据子文件夹名称自动应用标签？\n"
                        "（适用于按类别组织的文件夹结构）",
                        QMessageBox.Yes | QMessageBox.No,
                        QMessageBox.No
                    )
                    auto_label = reply == QMessageBox.Yes
                
                # 导入图像
                count = 0
                processed_count = 0
                
                # 显示进度对话框
                progress = QProgressDialog("正在导入图像...", "取消", 0, len(image_files), self)
                progress.setWindowModality(Qt.WindowModal)
                
                for i, img_path in enumerate(image_files):
                    # 更新进度
                    progress.setValue(i)
                    
                    # 检查是否取消
                    if progress.wasCanceled():
                        break
                    
                    # 导入图像
                    success = self.dataset_manager.add_image(img_path)
                    if success:
                        count += 1
                        
                        # 如果启用自动标注，根据文件夹名称应用标签
                        if auto_label:
                            # 获取图像所在的子文件夹名称
                            parent_folder = os.path.basename(os.path.dirname(img_path))
                            
                            # 如果标签不存在，创建它
                            if parent_folder and parent_folder not in [l["name"] for l in self.current_dataset["labels"]]:
                                # 为新标签选择随机颜色
                                color = QColor(
                                    np.random.randint(50, 200),
                                    np.random.randint(50, 200),
                                    np.random.randint(50, 200)
                                )
                                
                                # 添加新标签
                                self.dataset_manager.add_label(parent_folder, color.name())
                            
                            # 创建自动标注
                            if parent_folder:
                                # 获取图像尺寸
                                img = cv2.imread(img_path)
                                if img is not None:
                                    height, width = img.shape[:2]
                                    
                                    # 创建覆盖整个图像的标注
                                    annotation = [{
                                        "type": "rectangle",
                                        "x": 0,
                                        "y": 0,
                                        "width": width,
                                        "height": height,
                                        "label": parent_folder,
                                        "color": next((l["color"] for l in self.current_dataset["labels"] if l["name"] == parent_folder), "#FF0000")
                                    }]
                                    
                                    # 保存标注
                                    dataset_path = Path(self.current_dataset["path"])
                                    image_name = os.path.basename(img_path)
                                    annotation_file = dataset_path / "annotations" / f"{os.path.splitext(image_name)[0]}.json"
                                    
                                    # 确保annotations目录存在
                                    annotation_dir = annotation_file.parent
                                    if not annotation_dir.exists():
                                        annotation_dir.mkdir(parents=True, exist_ok=True)
                                    
                                    # 保存标注数据
                                    with open(annotation_file, "w", encoding="utf-8") as f:
                                        json.dump(annotation, f, ensure_ascii=False, indent=2)
                                    
                                    processed_count += 1
                
                    # 处理QApplication事件，保持UI响应
                    QApplication.processEvents()
                
                # 关闭进度对话框
                progress.setValue(len(image_files))
                
                # 更新界面
                self._update_image_list()
                self._update_labels()
                
                # 显示结果
                message = f"已成功导入{count}张图像到数据集"
                if auto_label and processed_count > 0:
                    message += f"，并为{processed_count}张图像创建了自动标注"
                
                QMessageBox.information(self, "导入成功", message)
                self.logger.info(message)
    
    def on_create_label(self):
        """创建标签响应函数"""
        self.logger.info("用户请求创建新标签")
        
        if not self.current_dataset:
            QMessageBox.warning(self, "操作失败", "请先打开或创建一个数据集")
            return
        
        # 显示添加标签对话框
        dialog = AddLabelDialog(self)
        result = dialog.exec_()
        
        if result == AddLabelDialog.Accepted:
            # 获取对话框结果
            label_name, label_color = dialog.get_result()
            
            # 验证标签名称
            if not label_name:
                QMessageBox.warning(self, "创建失败", "标签名称不能为空")
                return
            
            # 添加到数据集
            if self.dataset_manager.add_label(label_name, label_color.name()):
                # 更新标签管理器
                self._update_labels()
                self.logger.info(f"已创建标签: {label_name}")
            else:
                QMessageBox.warning(self, "创建失败", f"无法创建标签 '{label_name}'，可能是名称重复")
    
    def on_label_selected(self, label_name):
        """标签选择响应函数"""
        self.logger.debug(f"选择了标签: {label_name}")
        # 此处可以更新标注工具栏的当前标签
    
    def on_label_color_changed(self, label_name, color):
        """标签颜色变更响应函数"""
        self.logger.debug(f"标签 {label_name} 的颜色已更改为 {color.name()}")
        # 更新数据集中的标签颜色
        self.dataset_manager.update_label_color(label_name, color.name())
    
    def on_label_removed(self, label_name):
        """标签删除响应函数"""
        self.logger.debug(f"删除标签: {label_name}")
        # 从数据集中删除标签
        if self.dataset_manager.remove_label(label_name):
            self.logger.info(f"已从数据集移除标签: {label_name}")
        else:
            QMessageBox.warning(self, "删除失败", f"无法删除标签 '{label_name}'")
    
    def on_label_renamed(self, old_name, new_name):
        """标签重命名响应函数"""
        self.logger.debug(f"重命名标签: {old_name} -> {new_name}")
        # 更新数据集中的标签名称
        if self.dataset_manager.rename_label(old_name, new_name):
            self.logger.info(f"已重命名标签: {old_name} -> {new_name}")
        else:
            QMessageBox.warning(self, "重命名失败", f"无法重命名标签 '{old_name}'")
            # 重新加载标签列表以恢复原名称
            self._update_labels()
    
    def on_manage_labels(self):
        """管理标签响应函数"""
        self.logger.debug("用户请求管理标签")
        
        # 显示标签管理器面板
        self.label_dock.setVisible(True)
        self.label_dock.raise_()
    
    def _update_labels(self):
        """更新标签列表"""
        if self.current_dataset:
            labels = self.current_dataset.get("labels", [])
            self.label_manager.set_labels(labels)
    
    def _update_ui_for_dataset(self):
        """更新界面以适应当前数据集"""
        if self.current_dataset:
            # 隐藏欢迎文本
            self.welcome_label.setVisible(False)
            
            # 显示画布
            self.canvas.setVisible(True)
            
            # 显示停靠窗口
            self.dataset_dock.setVisible(True)
            self.label_dock.setVisible(True)
            self.training_dock.setVisible(True)
            
            # 启用相关动作
            self.import_images_action.setEnabled(True)
            self.import_folder_action.setEnabled(True)
            self.create_label_action.setEnabled(True)
            self.manage_labels_action.setEnabled(True)
            self.roi_action.setEnabled(True)
            
            # 更新标签管理器
            self._update_labels()
            
            # 更新图像列表
            self._update_image_list()
            
            # 如果数据集中有图像和标签，启用训练相关动作
            has_images = self.current_dataset.get("image_count", 0) > 0
            has_labels = len(self.current_dataset.get("labels", [])) > 0
            
            self.start_training_action.setEnabled(has_images and has_labels)
            self.export_model_action.setEnabled(has_images and has_labels)
            
            # 更新训练区域算法选择
            dataset_type = self.current_dataset.get("algorithm_type", "classification")
            self.training_monitor.filter_algorithms_by_type(dataset_type)
            
            # 设置ROI
            if self.current_dataset.get("roi"):
                roi_data = self.current_dataset.get("roi")
                roi = QRectF(
                    roi_data.get("x", 0), 
                    roi_data.get("y", 0),
                    roi_data.get("width", 1),
                    roi_data.get("height", 1)
                )
                self.canvas.set_roi(roi)
    
    def _update_image_list(self):
        """更新图像列表"""
        if self.current_dataset:
            # 获取数据集中的图像路径
            image_paths = self.dataset_manager.get_image_paths()
            
            # 更新数据集浏览器
            dataset_path = self.current_dataset.get("path", "")
            self.dataset_browser.set_dataset(dataset_path, image_paths)
            
            # 更新标注工具栏的标签列表
            labels = self.current_dataset.get("labels", [])
            self.annotation_toolbar_widget.set_classes(labels)
    
    def on_image_selected(self, image_path):
        """图像选择响应函数"""
        self.logger.debug(f"选择图像: {image_path}")
        
        # 在画布中显示图像
        self.canvas.load_image(image_path)
        
        # 加载图像的标注（如果有）
        self._load_annotations(image_path)
        
        # 启用标注工具栏
        self.annotation_toolbar.setVisible(True)
        
        # 保存当前选中的图像路径
        self.current_image_path = image_path
    
    def on_image_double_clicked(self, image_path):
        """图像双击响应函数"""
        self.logger.debug(f"双击图像: {image_path}")
        
        # 全屏显示图像或进入编辑模式
        # 可以在此处实现
    
    def on_image_deleted(self, image_path):
        """图像删除响应函数"""
        self.logger.debug(f"删除图像: {image_path}")
        
        try:
            # 删除图像文件
            os.remove(image_path)
            
            # 删除相关的标注文件（如果有）
            self._delete_annotations(image_path)
            
            # 更新数据集元数据
            if self.current_dataset:
                self.current_dataset["image_count"] -= 1
                self.dataset_manager._save_metadata()
                
            self.logger.info(f"已删除图像: {os.path.basename(image_path)}")
        except Exception as e:
            self.logger.error(f"删除图像失败: {e}")
            QMessageBox.warning(self, "删除失败", f"无法删除图像，错误: {e}")
    
    def _load_annotations(self, image_path):
        """加载图像的标注数据"""
        if not self.current_dataset or not image_path:
            return
        
        try:
            # 获取标注文件路径
            dataset_path = Path(self.current_dataset["path"])
            image_name = os.path.basename(image_path)
            annotation_file = dataset_path / "annotations" / f"{os.path.splitext(image_name)[0]}.json"
            
            if annotation_file.exists():
                # 读取标注数据
                with open(annotation_file, "r", encoding="utf-8") as f:
                    annotations = json.load(f)
                
                # 将标注加载到画布
                self.canvas.load_annotations(annotations)
                
                self.logger.debug(f"已加载标注: {annotation_file}")
            else:
                # 清除画布上的标注
                self.canvas.clear_annotations()
        except Exception as e:
            self.logger.error(f"加载标注失败: {e}")
            self.canvas.clear_annotations()
    
    def _delete_annotations(self, image_path):
        """删除图像的标注数据"""
        if not self.current_dataset or not image_path:
            return
        
        try:
            # 获取标注文件路径
            dataset_path = Path(self.current_dataset["path"])
            image_name = os.path.basename(image_path)
            annotation_file = dataset_path / "annotations" / f"{os.path.splitext(image_name)[0]}.json"
            
            if annotation_file.exists():
                # 删除标注文件
                os.remove(annotation_file)
                self.logger.debug(f"已删除标注文件: {annotation_file}")
        except Exception as e:
            self.logger.error(f"删除标注失败: {e}")
    
    def on_annotation_tool_changed(self, tool):
        """标注工具变更响应函数"""
        self.logger.debug(f"选择标注工具: {tool}")
        
        # 将工具选择传递给画布
        self.canvas.set_tool(tool)
    
    def on_annotation_color_changed(self, color):
        """标注颜色变更响应函数"""
        self.logger.debug(f"设置标注颜色: {color.name()}")
        
        # 将颜色选择传递给画布
        self.canvas.set_color(color)
    
    def on_annotation_class_changed(self, class_name):
        """标注类别变更响应函数"""
        self.logger.debug(f"设置标注类别: {class_name}")
        
        # 将类别选择传递给画布
        self.canvas.set_class(class_name)
    
    def on_start_training_action(self):
        """开始训练菜单动作响应函数"""
        self.logger.info("用户请求开始训练")
        
        # 显示训练监控面板
        self.training_dock.setVisible(True)
        self.training_dock.raise_()
    
    def on_start_training(self, params):
        """
        开始训练
        
        Args:
            params (dict): 训练参数
        """
        self.logger.info(f"开始训练模型，参数: {params}")
        
        if not self.current_dataset:
            QMessageBox.warning(self, "训练失败", "请先打开或创建一个数据集")
            return
        
        # 记录选择的算法
        if 'algorithm' in params:
            self.logger.info(f"选择的算法: {params['algorithm']}")
        
        # 实例化训练器
        self.trainer = ModelTrainer(
            self.current_dataset["path"],
            self.current_dataset["algorithm_type"],
            params,
            self.logger
        )
        
        # 设置回调
        self.trainer.on_progress = self._on_training_progress
        self.trainer.on_completed = self._on_training_completed
        
        # 启动训练线程
        self.trainer.start()
    
    def on_stop_training(self):
        """停止训练响应函数"""
        self.logger.info("停止训练模型")
        
        if hasattr(self, 'trainer') and self.trainer.running:
            self.trainer.stop()
    
    def _on_training_progress(self, epoch, metrics):
        """训练进度回调函数"""
        # 在主线程中更新UI
        self.training_monitor.update_progress(epoch, metrics)
    
    def _on_training_completed(self, success, metrics):
        """训练完成回调函数"""
        self.logger.info(f"训练{'成功' if success else '失败'}")
        
        # 在主线程中更新UI
        self.training_monitor.training_completed(metrics)
        
        # 如果成功，启用导出模型功能
        if success:
            self.export_model_action.setEnabled(True)
    
    def on_export_model(self):
        """导出模型响应函数"""
        self.logger.info("用户请求导出模型")
        
        if not self.current_dataset:
            QMessageBox.warning(self, "导出失败", "请先打开或创建一个数据集")
            return
        
        if not hasattr(self, 'trainer') or not self.trainer.model:
            QMessageBox.warning(self, "导出失败", "请先训练模型")
            return
        
        # 打开保存对话框
        file_dialog = QFileDialog(self)
        file_dialog.setAcceptMode(QFileDialog.AcceptSave)
        file_dialog.setNameFilter("模型文件 (*.onnx *.pt *.pkl *.h5)")
        file_dialog.setDefaultSuffix("onnx")
        
        if file_dialog.exec_():
            file_path = file_dialog.selectedFiles()[0]
            if file_path:
                # 导出模型
                if self.trainer.export_model(file_path):
                    QMessageBox.information(self, "导出成功", f"模型已导出到: {file_path}")
                else:
                    QMessageBox.warning(self, "导出失败", "模型导出失败，请查看日志")
    
    def save_annotations(self):
        """保存当前图像的标注"""
        if not self.current_dataset or not hasattr(self, 'current_image_path') or not self.current_image_path:
            return False
        
        try:
            # 获取标注数据
            annotations = self.canvas.get_annotations()
            
            # 如果没有标注，则跳过
            if not annotations:
                return False
            
            # 获取标注文件路径
            dataset_path = Path(self.current_dataset["path"])
            image_name = os.path.basename(self.current_image_path)
            annotation_file = dataset_path / "annotations" / f"{os.path.splitext(image_name)[0]}.json"
            
            # 确保annotations目录存在
            annotation_dir = annotation_file.parent
            if not annotation_dir.exists():
                annotation_dir.mkdir(parents=True, exist_ok=True)
            
            # 保存标注数据
            with open(annotation_file, "w", encoding="utf-8") as f:
                json.dump(annotations, f, ensure_ascii=False, indent=2)
            
            # 更新数据集标注计数
            self.current_dataset["annotation_count"] = len(os.listdir(dataset_path / "annotations"))
            self.dataset_manager._save_metadata()
            
            self.logger.debug(f"已保存标注: {annotation_file}")
            return True
        except Exception as e:
            self.logger.error(f"保存标注失败: {e}")
            return False
    
    def on_annotation_added(self, annotation):
        """标注添加响应函数"""
        self.logger.debug(f"添加标注: {annotation}")
        
        # 保存标注
        self.save_annotations()
        
        # 更新标签计数
        if annotation.get('label'):
            label_name = annotation.get('label')
            for label in self.current_dataset["labels"]:
                if label["name"] == label_name:
                    label["count"] += 1
                    break
    
    def on_annotation_deleted(self, annotation_index):
        """标注删除响应函数"""
        self.logger.debug(f"删除标注: 索引 {annotation_index}")
        
        # 保存标注变更
        self.save_annotations()
        
        # 更新标签计数
        # 由于标注已被删除，此处我们只需保存当前状态即可
        self.dataset_manager._save_metadata()
    
    def on_set_roi(self):
        """设置ROI区域响应函数"""
        self.logger.info("用户请求设置ROI区域")
        
        if not self.current_dataset:
            QMessageBox.warning(self, "操作失败", "请先打开或创建一个数据集")
            return
        
        # 获取当前图像
        if not hasattr(self, 'current_image_path') or not self.current_image_path:
            # 尝试获取数据集的第一张图像
            image_paths = self.dataset_manager.get_image_paths()
            if not image_paths:
                QMessageBox.warning(self, "操作失败", "数据集中没有图像")
                return
            current_image = cv2.imread(image_paths[0])
        else:
            current_image = cv2.imread(self.current_image_path)
        
        if current_image is None:
            QMessageBox.warning(self, "操作失败", "无法加载图像")
            return
        
        # 获取当前ROI设置
        current_roi = None
        if self.current_dataset.get("roi"):
            roi_data = self.current_dataset.get("roi")
            current_roi = QRectF(
                roi_data.get("x", 0.25),
                roi_data.get("y", 0.25),
                roi_data.get("width", 0.5),
                roi_data.get("height", 0.5)
            )
        
        # 创建并显示ROI设置对话框
        dialog = ROIDialog(current_image, current_roi, self)
        result = dialog.exec_()
        
        if result == ROIDialog.Accepted:
            # 获取对话框结果
            roi, apply_to_all = dialog.get_roi()
            
            # 保存ROI设置
            self.current_dataset["roi"] = {
                "x": roi.x(),
                "y": roi.y(),
                "width": roi.width(),
                "height": roi.height()
            }
            
            # 保存到数据集元数据
            self.dataset_manager._save_metadata()
            
            # 应用到当前图像
            self.canvas.set_roi(roi)
            
            self.logger.info("已设置ROI区域")
            
            # 显示确认消息
            if apply_to_all:
                QMessageBox.information(self, "设置成功", "ROI设置已应用于所有图像")
            else:
                QMessageBox.information(self, "设置成功", "ROI设置已应用于当前图像")
    
    def on_convert_to_annotation(self, true_class, pred_class, sample_ids):
        """
        将错误分类的样本转换为标注任务
        
        Args:
            true_class (str): 真实类别
            pred_class (str): 预测类别
            sample_ids (list): 样本ID列表
        """
        self.logger.info(f"准备将{len(sample_ids)}个被错误分类为[{pred_class}]的[{true_class}]类样本转换为标注任务")
        
        # 检查当前数据集
        if not self.current_dataset:
            QMessageBox.warning(self, "转换失败", "没有打开的数据集")
            return
        
        # 确认对话框
        reply = QMessageBox.question(
            self,
            "转换为标注任务",
            f"将把{len(sample_ids)}个被错误分类为[{pred_class}]的[{true_class}]类样本转换为标注任务。\n"
            "这将重置这些样本的标注。是否继续？",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No
        )
        
        if reply != QMessageBox.Yes:
            return
        
        # 获取样本图像路径
        dataset_images = self.dataset_manager.get_images()
        
        # 创建批量处理进度对话框
        progress = QProgressDialog("正在处理样本...", "取消", 0, len(sample_ids), self)
        progress.setWindowModality(Qt.WindowModal)
        
        # 计数器
        processed_count = 0
        
        # 处理每个样本
        for i, sample_id in enumerate(sample_ids):
            # 更新进度
            progress.setValue(i)
            
            # 检查取消
            if progress.wasCanceled():
                break
            
            # 查找样本图像
            image_path = None
            for img in dataset_images:
                # 样本ID可能是完整路径或只是文件名
                if sample_id == img or os.path.basename(img) == sample_id:
                    image_path = img
                    break
            
            if not image_path:
                continue
            
            # 删除现有标注
            self._delete_annotations(image_path)
            
            # 创建新的矩形标注，覆盖整个图像
            try:
                # 读取图像获取尺寸
                img = cv2.imread(image_path)
                if img is None:
                    continue
                    
                height, width = img.shape[:2]
                
                # 获取正确标签的颜色
                label_color = "#FF0000"  # 默认红色
                for label in self.current_dataset["labels"]:
                    if label["name"] == true_class:
                        label_color = label["color"]
                        break
                
                # 创建覆盖整个图像的矩形标注
                annotation = [{
                    "type": "rectangle",
                    "x": 0,
                    "y": 0,
                    "width": width,
                    "height": height,
                    "label": true_class,
                    "color": label_color
                }]
                
                # 保存标注
                dataset_path = Path(self.current_dataset["path"])
                image_name = os.path.basename(image_path)
                annotation_file = dataset_path / "annotations" / f"{os.path.splitext(image_name)[0]}.json"
                
                # 确保annotations目录存在
                annotation_dir = annotation_file.parent
                if not annotation_dir.exists():
                    annotation_dir.mkdir(parents=True, exist_ok=True)
                
                # 保存标注数据
                with open(annotation_file, "w", encoding="utf-8") as f:
                    json.dump(annotation, f, ensure_ascii=False, indent=2)
                
                processed_count += 1
                
            except Exception as e:
                self.logger.error(f"处理样本 {sample_id} 时出错: {e}")
                continue
            
            # 处理QApplication事件，保持UI响应
            QApplication.processEvents()
        
        # 关闭进度对话框
        progress.setValue(len(sample_ids))
        
        # 更新UI
        self._update_image_list()
        
        # 显示结果
        QMessageBox.information(
            self,
            "转换完成",
            f"已将{processed_count}个样本转换为标注任务。您现在可以在数据集浏览器中查看这些样本并调整标注。"
        )
        
        self.logger.info(f"成功将{processed_count}个样本转换为标注任务")
    
    def closeEvent(self, event):
        """窗口关闭事件"""
        # 询问用户是否确定退出
        reply = QMessageBox.question(
            self, 
            "确认退出", 
            "确定要退出程序吗？",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            self.logger.info("应用程序关闭")
            event.accept()
        else:
            event.ignore() 